package com.example.websocket.service;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author ZhenWuWang
 */
@Component
public class SpringWebSocketHandler extends TextWebSocketHandler
{
    private final ConcurrentHashMap<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();



    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception
    {
        System.out.println("connect to the websocket success......当前数量:" + sessionMap.size());
        String userName = (String) session.getAttributes().get("WEBSOCKET_USERNAME");
        sessionMap.put(userName, session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception
    {
        String userName = (String) session.getAttributes().get("WEBSOCKET_USERNAME");
        System.out.println("user:" + userName + "退出登陆");
        sessionMap.remove(userName);
    }

    /**
     * js调用websocket.send时候，会调用该方法
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception
    {
        String userName = (String) session.getAttributes().get("WEBSOCKET_USERNAME");
        System.out.println("接收到用户" + userName + "发送过来的消息：" + message.toString());
        session.sendMessage(new TextMessage("replay msg: " + message.getPayload()));
        sendMessageToAllUser(message);
        session.sendMessage(new TextMessage("over"));
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception
    {
        if (session.isOpen())
        {
            session.close();
        }
        String userName = (String) session.getAttributes().get("WEBSOCKET_USERNAME");
        System.out.println("用户" + userName + "的session已经断开...");
        sessionMap.remove(userName);
    }

    @Override
    public boolean supportsPartialMessages()
    {
        return false;
    }

    public void sendMessageToUser(String userName, TextMessage message)
    {
        WebSocketSession session = sessionMap.get(userName);
        if (session != null)
        {
            try
            {
                if (session.isOpen())
                {
                    session.sendMessage(message);
                }
            } catch (IOException e)
            {
                e.printStackTrace();
            }
        }
    }

    public void sendMessageToAllUser(TextMessage message)
    {
        sessionMap.forEach((userName, session) ->
        {
            if (session != null && session.isOpen())
            {
                try
                {
                    session.sendMessage(message);
                } catch (IOException e)
                {
                    e.printStackTrace();
                }
            }
        });
    }
}
